Fuse online softmax: compute running max and sum in one tree reduction - #3
Conversation
|
Hi there, thanks for the contribution! Looks good to me, and I love that you actually implemented the proper online softmax. Please include add some tests or at least show some proofs that it works to make sure there is no regression, and then we can merge. Thank you! |
|
Hi, are you going to update the PR or should I close it? Thanks! |
|
Hi, sorry about the delay. I missed your earlier request to update the PR. I’d still like to finish it. Would it be okay if I take two days to resolve the conflicts and add the requested tests/validation? I’ll update the PR within two days. Thanks for following up! |
|
@pranathi000 Sure, no rush |
Added detailed comments explaining the correctness of the softmax fusion kernel.
|
Sorry for the delay in replying. You asked for tests or proof that this doesn't break anything, so here's what I did, in three parts.
I took the kernel from before my change and the fused version from this branch, gave both the exact same input row, and printed out what each one produced. Every value came out bit-for-bit identical between the two. This tells me the fusion didn't change the actual output for a normal case.
One row won't catch every possible bug, so I also ran both kernels on a row where every value is tied for the max, a row with one very large value sitting next to a bunch of tiny ones to check the numerical stability part specifically, a row of all negative numbers, and a bigger row size than the first test. Old and new matched exactly on every one of these as well, no missing values, no NaN, no infinity anywhere.
Testing only proves it works for the specific inputs you tried, so I also worked through the math behind the merge step, to show it holds for any input, not just the ones I picked. I wrote that reasoning out as a comment at the top of src/test_softmax.cu, along with the real test code for both steps above, so you can read through the logic and rerun the tests yourself if you want to check it. Let me know if you need anything else before merging. |
|
Quick note on what src/test_softmax.cu is actually doing. It’s a small standalone test file, not part of the main engine. It builds a few test rows, runs the real I went with this instead of just pasting printed numbers because you can actually run the test and get a pass/fail result, rather than having to take my word for it. I also didn’t pick arbitrary rows. One row has all values tied for the maximum, to check the case where there is no single winner. One row has a very large value next to a bunch of tiny ones, specifically to exercise the numerical-stability part, since that is the reason for subtracting the maximum in the first place. Then there’s a row containing all negative values and a larger row size, just to make sure the same properties hold outside the simplest case. The comment at the top of the same file also walks through why the merge step is mathematically equivalent to the old two-pass version, rather than only showing that it happened to work for these particular test inputs. |
|
Nice, please pull the test file to |
This file implements a correctness check for the fused online softmax kernel, testing various edge cases.
|
Moved the test file to tests/test_softmax.cu. Reran it there and all 5 cases still pass. Also built the full project through CMake to double check, not just my test file. Ran into two things along the way, both on my end, not changes I made to the repo. CMAKE_CUDA_COMPILER is set to /opt/cuda/bin/nvcc, but on Colab nvcc is at /usr/local/cuda/bin/nvcc instead. That path is set directly in the file, so I couldn't override it from the command line, I edited that one line locally just to point it at the right path for Colab. Didn't commit or push that change anywhere. CMAKE_CUDA_ARCHITECTURES is also set to 120 in the file, which is a newer GPU generation than what I have. I passed -DCMAKE_CUDA_ARCHITECTURES=75 on the command line to match my GPU instead. After both of those, the build went through cleanly, tiny-vllm compiled and linked with no errors. Both of these look like they're just set to match your own local setup, so wanted to flag it in case it's useful, but didn't want to change anything in the PR itself since it's unrelated to the softmax fix. Let me know if you need anything else before merging. |
2 participants
This change fuses the two softmax reductions into one.
Currently both softmaxKernel and softmaxKernelDecode do two full tree reductions per row: the first finds the maximum, and the second sums the exponentials, each with its own __syncthreads() loop. This PR combines them into a single reduction using the online softmax formulation already used per-token in pagedAttentionKernel.
Each thread now carries a (max, sum) pair up the tree and merges them together at every level:
m = max(m_a, m_b)
d = d_a * exp(m_a - m) + d_b * exp(m_b - m)
The exp(m_a - m) term rescales a partial sum whenever a larger max appears, which keeps the computation numerically stable. When the reduction reaches the root, both the row max and the normalized denominator are already available, so the separate max broadcast and the intermediate exponentiation pass are no longer needed.
This roughly halves the number of __syncthreads() barriers per row. It adds a couple of expf calls per merge, but on GPU the reduction in barriers is the better trade-off.
Kernel signatures, launch configurations, and the existing 1024-thread guard are unchanged. Shared memory usage goes from a single float[1024] (plus one scalar) to two float[1024] arrays.